from netmiko import ConnectHandler


# 添加ip白名单
def add_ip(squid_info, hosts, description):
    net_connect = ConnectHandler(**squid_info)
    print("已登录服务器：" + squid_info["host"])
    # 初始化命令列表
    commands = []
    # 循环读取逗号分割的host
    for host in hosts.split(","):
        commands.append('echo acl localnet src ' + host + '\ \ \ \ ' + '#' + description + ' >> /etc/squid/srcip.acl')
    commands.append('systemctl restart squid')
    print("执行命令如下：")
    # for command in commands:
    #     print(command)
    # 执行
    print(net_connect.find_prompt(), end=' ')
    result = net_connect.send_config_set(commands)
    print(result)

    # 关闭连接
    net_connect.disconnect()


# 删除ip白名单
def del_ip(squid_info, hosts):
    net_connect = ConnectHandler(**squid_info)
    print("已登录服务器：" + squid_info["host"])
    # 初始化命令列表
    commands = []
    # 循环读取逗号分割的host
    for host in hosts.split(","):
        commands.append("sed -i '/" + host + "/d' /etc/squid/ipacl.conf")
    commands.append('systemctl restart squid')
    print("执行命令如下：")
    # for command in commands:
    #     print(command)
    # 执行
    print(net_connect.find_prompt(), end=' ')
    result = net_connect.send_config_set(commands)
    print(result)

    # 关闭连接
    net_connect.disconnect()


# 添加域名白名单
def add_domain(squid_info, dstdomains):
    net_connect = ConnectHandler(**squid_info)
    print("已登录服务器：" + squid_info["host"])

    # 初始化命令列表
    commands = []
    # 循环读取逗号分割的host

    for dstdomain in dstdomains.split(","):
        commands.append('echo acl White_List dstdomain ' + dstdomain + ' >> /etc/squid/desdomain.acl')
    commands.append('systemctl restart squid')
    print("执行命令如下：")
    # for command in commands:
    #     print(command)
    # 执行
    print(net_connect.find_prompt(), end=' ')
    result = net_connect.send_config_set(commands)
    print(result)
    # 关闭连接
    net_connect.disconnect()


def main():
    # 定义连接信息
    all_squid_info = [
        {
            'device_type': 'linux',
            'host': '172.20.68.9',
            'username': 'root',
            'password': '5Star.com'
        },
        {
            'device_type': 'linux',
            'host': '172.20.68.10',
            'username': 'root',
            'password': '5Star.com'
        }
    ]

    print("***squid代理服务器变更脚本***")
    print("---1、批量添加ip")
    print("---2、批量添加域名白名单")

    choice = int(input("请输入选择："))
    if choice == 1:
        hosts = input("请输入ip(批量请用英文逗号隔开）:")
        description = input("请输入描述：")
        for squid_info in all_squid_info:
            add_ip(squid_info, hosts, description)
            print("squid重启完毕")
            print("===================\n")

    elif choice == 2:
        dstdomain = input("请输入域名(批量请用英文逗号隔开）:")
        for squid_info in all_squid_info:
            add_domain(squid_info, dstdomain)
            print("squid重启完毕")
            print("===================\n")

    elif choice == 3:
        hosts = input("请输入ip(批量请用英文逗号隔开）:")
        for squid_info in all_squid_info:
            del_ip(squid_info, hosts)
            print("squid重启完毕")
            print("===================\n")
    pass


if __name__ == '__main__':
    main()
